#!/usr/bin/python3
"""
单线程去解析网络数据包。
"""

from scapy.all import *
import re
import os
from threading import Thread

file_size = 1 * 1000 * 1000
dir_path = '/root/tcp/'
save_path = dir_path + 'statistic.txt'


class Parser:
    def __init__(self):
        self.ip = {}
        self.dir = []

    def parse_ip(self):
        while 1:
            if self.dir:
                for _dir in self.dir:
                    start = time.time()
                    packets = rdpcap(filename=_dir)
                    for p in packets:
                        src_dst = packets._elt2sum(p).split('>')
                        for ip in src_dst:
                            ip_search = re.search(r'[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+', ip)
                            if ip_search is None:
                                continue
                            _ip = ip[ip_search.start():ip_search.end()]
                            if _ip in self.ip.keys():
                                self.ip[_ip] = self.ip[_ip] + 1
                            else:
                                self.ip[_ip] = 1
                    self.dir.remove(_dir)
                    os.remove(_dir)
                    print('{} {} {}'.format('del', time.time() - start, len(self.dir)))

    def add_pcap(self):
        while 1:
            list_dir = os.listdir(dir_path)
            for _dir in list_dir:
                file_path = dir_path + _dir
                if os.path.isfile(file_path):
                    if os.path.getsize(file_path) > file_size:
                        if file_path not in self.dir:
                            self.dir.append(file_path)
            time.sleep(2)

    def save_file(self):
        while 1:
            with open(save_path, 'w')as writer:
                for key in self.ip.keys():
                    writer.write('{:<20}{:>20}\n'.format(key, self.ip[key]))
            time.sleep(2)


if __name__ == '__main__':
    p = Parser()
    t1 = Thread(target=p.parse_ip)
    t2 = Thread(target=p.add_pcap)
    t3 = Thread(target=p.save_file)
    t1.start()
    t2.start()
    t3.start()
